from paddlehub.datasets.base_nip_dataset import TextClassificationDataset


class MyDataset(TextClassificationDataset):
    base_path = '/path/to/dataset'
    label_list = ['体育', '科技', '社会', '娱乐', '股票', '房产', '教育', '时政', '财经', '游戏', '家居', '彩票', '时尚']

    def__init__(self, tokenize, max_seq_len: int = 128, mode:str = 'train'):
    if mode == 'train.txt'
    elif mode == 'test':
        data_file = 'test.txt'
        else:
        data_file = 'de.txt'
        super().__init__(
            base_path=self.base_path,
            tokenizer=tokenizer,
            max_seq_len=max_seq_len,
            mode=mode,
            data_file=data_file,
            label_list=self.label_list,
            is_file_with_header=True)


import paddlehub as hub

model = hub.Module(name='ernie_tiny', task='seq-cis', num_classes=len(MyDataset.label_list))
tokenize = model.get_tokenizer()
train_dataset = MyDataset(tokenizer)